import os
from zhipuai import ZhipuAI
from dotenv import load_dotenv, find_dotenv

_ = load_dotenv(find_dotenv())
api_key = os.environ.get('ZHIPUAI_API_KEY')
if api_key is None:
    raise ValueError("API Key is not set in the .env file")
client = ZhipuAI(api_key=api_key)

def get_completion(prompt, model="glm-4-plus", temperature=0.01):
    messages = [{"role": "user", "content": prompt}]
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=temperature  
    )
    return response.choices[0].message.content
  
import streamlit as st

# 头像配置
ICON_AI = '💻'
ICON_USER = '🧑'

# 显示一条消息（包含头像与消息内容）
def dspMessage(role, content):
    with st.chat_message(role, 
                         avatar=ICON_AI if role == 'assistant' else ICON_USER):
        st.write(content)

# 追加并显示一条消息
def append_and_show(role, content):
    st.session_state.messages.append({"role": role, "content": content})      
    dspMessage(role, content)

# 如果还没有消息，则添加第一条提示消息
if 'messages' not in st.session_state:
    st.session_state.messages = [{"role": "assistant", 
                                  "content": "我是你的信息助手，请问你查询什么信息？"}]

# 将会话中的messages列表中的消息全部显示出来
for msg in st.session_state.messages:
    dspMessage(msg["role"], msg["content"])

# 接受用户输入的提示词，并调用大模型API获得反馈
if prompt := st.chat_input():
    append_and_show("user", prompt)
    reply = get_completion(prompt) 
    append_and_show("assistant", reply) 